# run_pso_multiprocess.py (单一高精度模式 - 最终版)
# 描述: 直接使用高精度评估模型，执行单一阶段的PSO优化。

import pickle
import time
import numpy as np
import multiprocessing as mp
import logging
# 假设您的核心类都在 pso_core.py 文件中
from pso_core import HierarchicalRoadNetwork, OptimizedMultiSitePSO

# (setup_progress_logger 和 progress_callback 函数保持不变)
def setup_progress_logger():
    # ... (此函数代码保持不变) ...
    progress_logger = logging.getLogger('pso_progress')
    progress_logger.setLevel(logging.INFO)
    if progress_logger.hasHandlers():
        progress_logger.handlers.clear()
    fh = logging.FileHandler('pso_progress.log', mode='w')
    formatter = logging.Formatter('%(message)s')
    fh.setFormatter(formatter)
    progress_logger.addHandler(fh)
    return progress_logger

def progress_callback(iteration, max_iter, current_best_score):
    # ... (此函数代码保持不变) ...
    progress_logger = logging.getLogger('pso_progress')
    progress_logger.info(f"Iteration {iteration}/{max_iter} | Best Score: {current_best_score:.4f}")
    for handler in progress_logger.handlers:
        handler.flush()

if __name__ == "__main__":
    mp.set_start_method('spawn', force=True)
    setup_progress_logger()

    print("\n--- 多进程优化脚本已启动 (单一高精度模式) ---")
    
    # --- 1. 加载共享数据 ---
    print("正在加载输入数据...")
    with open("pso_input_data.pkl", "rb") as f:
        pso_input_data = pickle.load(f)
    boundary_polygon_3857 = pso_input_data["boundary_polygon_3857"]
    poi_coords_3857 = pso_input_data["poi_coords_3857"]
    population_points_3857 = pso_input_data["population_points_3857"]
    road_network_cache = pso_input_data["road_network_cache"]
    road_network = HierarchicalRoadNetwork(road_gdf=None, cache_file=road_network_cache)

    print("\n" + "="*20 + " 开始执行高精度优化 (Precise Mode) " + "="*20)
    
    # 定义PSO参数
    pso_kwargs = {
        'poi_coords': poi_coords_3857, 
        'population_points': population_points_3857,
        'num_particles': 20,       # 可以根据需要调整粒子数
        'max_iter': 100,           # 迭代次数可以适当增加，因为没有粗搜阶段
        'radius': 3000, 
        'max_search_radius': 3000,
        'poi_weight': 0.7, 
        'pop_weight': 0.2, 
        'spacing_weight': 0.1,
        'early_stop_patience': 30, # 提前停止的耐心值也可以调整
        'use_adaptive_params': True # 仍然推荐使用动态PSO参数
    }
    pso_optimizer = OptimizedMultiSitePSO(
        num_sites=30, 
        boundary_polygon=boundary_polygon_3857,
        road_network=road_network, 
        **pso_kwargs
    )
    
    # 🔥 直接调用 fit 方法，并明确指定 evaluation_mode='precise'
    final_locations = pso_optimizer.fit(
        progress_callback=progress_callback,
        evaluation_mode='precise' 
    )

    # =========================================================================
    # === 结果保存 ============================================================
    # =========================================================================
    if final_locations:
        print("\n优化成功！正在进行最终方案分析...")
        final_position_3857 = np.array(final_locations).flatten()
        analysis_metrics = pso_optimizer.analyze_solution(final_position_3857)
        
        # 现在收敛历史只有一个阶段，直接保存即可
        convergence_history = pso_optimizer.g_best_scores

        results_to_save = {
            "optimal_locations_3857": final_locations,
            "analysis_metrics": analysis_metrics,
            "convergence_history": convergence_history,
            # 不再需要 stage1_end_iteration
        }

        # 在保存前做最后一次检查
        if not convergence_history or isinstance(convergence_history[0], (int, float)):
            with open("pso_results.pkl", "wb") as f:
                pickle.dump(results_to_save, f)
            print("✅ 优化结果已成功保存到 'pso_results.pkl'")
        else:
            print("❌ 保存失败：收敛历史数据格式不正确！")
    else:
        print("❌ 优化失败，未生成结果文件。")

    print("\n--- 多进程优化脚本执行完毕 ---")